package com.yueny.study.db.kernel.join.hash;

import cn.hutool.core.lang.Assert;
import cn.hutool.core.util.RandomUtil;
import com.yueny.study.db.kernel.join.data.Expr;
import com.yueny.study.db.kernel.join.data.PrepareData;
import com.yueny.study.db.kernel.join.data.Types;
import com.yueny.study.db.kernel.join.data.column.Column;
import com.yueny.study.db.kernel.join.data.row.JoinKey;
import com.yueny.study.db.kernel.join.data.row.Row;
import com.yueny.study.db.kernel.join.data.value.IntValue;
import com.yueny.study.db.kernel.join.data.value.StringValue;
import com.yueny.study.db.kernel.join.data.value.Value;

import java.util.AbstractMap;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author fengyang
 * @date 2023/10/27 下午3:08
 * @see <a href="yueny09@163.com">mailTo: yueny09@163.com</a>
 */
public class SimpleHashJoin implements IHashJoin
{
    private static int toHash(String data) {
        return data.hashCode();
    }

    private boolean isPrintDetailLog;

    public SimpleHashJoin()
    {
        this(false);
    }

    public SimpleHashJoin(boolean isPrintDetailLog)
    {
        this.isPrintDetailLog = isPrintDetailLog;
    }

    @Override
    public PrepareData prepareData()
    {
        int mockLeftRows = 2000000;
        int mockRightRows = 10;

        Column leftPKColumn = Column.builder()
                .columnType(Types.INT)
                .columnIdx(0)
                .columnName("id")
                .build();
        Column leftNameColumn = Column.builder()
                .columnType(Types.STRING)
                .columnIdx(1)
                .columnName("name")
                .build();
        Column leftAgeColumn = Column.builder()
                .columnType(Types.INT)
                .columnIdx(2)
                .columnName("age")
                .build();
        Column leftDescColumn = Column.builder()
                .columnType(Types.STRING)
                .columnIdx(3)
                .columnName("desc")
                .build();

        Column rightDescColumn = Column.builder()
                .columnType(Types.STRING)
                .columnIdx(0)
                .columnName("desc")
                .build();
        Column rightPKColumn = Column.builder()
                .columnType(Types.INT)
                .columnIdx(1)
                .columnName("id")
                .build();

        // step 1. 初始化表数据
        // 左表 2000行数据
        List<Row> leftRows = new ArrayList<>(mockLeftRows);
        for (int i = 0; i < mockLeftRows; i++) {
            // 假设左表有三列，分别为 id, name, age
            Value pkVal = new IntValue(i);
            Value nameVal = new StringValue(RandomUtil.randomString(5));
            Value ageVal = new IntValue(RandomUtil.randomInt(60));
            Value descVal = new StringValue(RandomUtil.randomString(2) + "-" + pkVal);

            Row r = Row.builder()
                    .columns(Arrays.asList(leftPKColumn, leftNameColumn, leftAgeColumn, leftDescColumn))
                    .values(Arrays.asList(pkVal, nameVal, ageVal, descVal))
                    .pkColumnList(Arrays.asList(leftPKColumn.getColumnName()))
                    .build();
            leftRows.add(r);
        }

        // 右表 10行数据
        List<Row> rightRows = new ArrayList<>(mockRightRows);
        for (int i = 0; i < mockRightRows; i++) {
            // 假设右表有两列，分别为 id, city
            Value pkVal = new IntValue(RandomUtil.randomInt(mockLeftRows));
            Value descVal = new StringValue(RandomUtil.randomString(2) + "-" + pkVal);

            Row r = Row.builder()
                    .columns(Arrays.asList(rightDescColumn, rightPKColumn))
                    .values(Arrays.asList(descVal, pkVal))
                    .pkColumnList(Arrays.asList(rightPKColumn.getColumnName()))
                    .build();
            rightRows.add(r);
        }

        JoinKey joinKey = JoinKey.builder()
                .expr(Expr.EQ)
                .left(leftPKColumn)
                .right(rightPKColumn)
                .build();

        return PrepareData.builder()
                .leftRows(leftRows)
                .rightRows(rightRows)
                .mockLeftRows(mockLeftRows).mockRightRows(mockRightRows)
                .leftPKColumn(leftPKColumn).rightPKColumn(rightPKColumn)
                .joinKey(joinKey)
                .build();
    }

    @Override
    public Map.Entry<List<Row>, List<Row>> assemblyTable(PrepareData prepareData)
    {
        // 选择 Build 和 Probe
        List<Row> build;
        List<Row> probe;
        if(prepareData.getLeftRows().size() < prepareData.getRightRows().size()) {
            build = prepareData.getLeftRows();
            probe = prepareData.getRightRows();
        } else {
            build = prepareData.getRightRows();
            probe = prepareData.getLeftRows();
        }

        return new AbstractMap.SimpleEntry<>(build, probe);
    }

    @Override
    public Map<Integer, Row> buildHash(List<Row> build)
    {
        // 将build中的每一条记录的主键 key 值计算为一个 hash 值，然后根据 hash 值将该记录插入到hash 表中
        Map<Integer /* hash value*/, Row /* row data */> buildHashTable = new HashMap<>();
        for (Row buildRow : build) {
            // get pk
            final List<String> pkColumnNameList = buildRow.getPkColumnList();
            // get pk 值， 字符串拼接，计算Hash
            StringBuffer sb = new StringBuffer();
            List<Column> pkColumnList = buildRow.getColumns().stream().filter(column -> {
                return pkColumnNameList.contains(column.getColumnName());
            }).collect(Collectors.toList());
            for (Column pkColumn : pkColumnList) {
                sb.append(buildRow.getValues().get(buildRow.getColumns().indexOf(pkColumn)).get());
            }

            buildHashTable.put(toHash(sb.toString()), buildRow);
        }

        return buildHashTable;
    }

    @Override
    public List<HasherMatchResult> hasherMatcher(List<Row> probe, final Map<Integer /* hash value*/, Row /* row data */> buildHashTable)
    {
        List<HasherMatchResult> matchResult = new ArrayList<>();
        // 使用和 build 中相同的 hash 函数，计算出相应的 hash 值，
        // 然后根据 hash 值在 hash 表中寻找到需要比较的记录，一一比较，得到最终结果。
        for (Row probeRow : probe) {
            // get pk
            List<String> probePkColumnNameList = probeRow.getPkColumnList();
            // get pk 值， 字符串拼接，计算Hash
            StringBuffer probePKHashBuffer = new StringBuffer();
            List<Column> pkColumnList = probeRow.getColumns().stream().filter(column -> {
                return probePkColumnNameList.contains(column.getColumnName());
            }).collect(Collectors.toList());
            for (Column pkColumn : pkColumnList) {
                probePKHashBuffer.append(probeRow.getValues().get(probeRow.getColumns().indexOf(pkColumn)).get());
            }
            int probePKHash = toHash(probePKHashBuffer.toString());

            Row buildRow = buildHashTable.get(probePKHash); // 在哈希表中查找匹配记录
            if (buildRow != null) {
                matchResult.add(
                        HasherMatchResult.builder()
                                .buildRow(buildRow).probeRow(probeRow)
                                .build()
                );
            }
        }

        return matchResult;
    }

    public static void main(String[] args) {
        SimpleHashJoin join = new SimpleHashJoin();

        PrepareData fullData = join.prepareData();
        Assert.isTrue(fullData.getLeftRows().size() == fullData.getMockLeftRows());
        Assert.isTrue(fullData.getRightRows().size() == fullData.getMockRightRows());

        // 耗时统计 start
        // 选择 Build 和 Probe
        Map.Entry<List<Row>, List<Row>> assemblyTable = join.assemblyTable(fullData);

        List<Row> build = assemblyTable.getKey();
        List<Row> probe = assemblyTable.getValue();

        // step 2. 根据 build 表构造Hash表
        Map<Integer /* hash value*/, Row /* row data */> buildHashTable = join.buildHash(build);

        // step 3. 进行哈希连接
        List<HasherMatchResult> hasherMatcher = join.hasherMatcher(probe, buildHashTable);
        // 耗时统计 end

        if(join.isPrintDetailLog) {
            // 输出
            System.out.println("Match size: " + hasherMatcher.size());
            for (HasherMatchResult item : hasherMatcher) {
                System.out.println("PK: " + item.getBuildRow().getPkColumnList() +
                        ", data = " + item.getBuildRow().getValues() + " join " + item.getProbeRow().getValues());
            }
        }
    }
}
